import os
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
import os
import os.path as osp
ROOT = osp.dirname(osp.abspath(__file__))

ENV_PATH = os.environ['CONDA_PREFIX']
if os.name == 'nt':
    ENV_PATH = os.path.join(ENV_PATH, 'Library')

if os.name == 'nt':
    include_dirs = [os.path.join(ENV_PATH, 'include'), '.']
else:
    include_dirs = [os.path.join(ENV_PATH, 'include', 'opencv4'), '.']

library_dirs = [os.path.join(ENV_PATH, 'lib')]

setup(
    name='gaussianpro',
    ext_modules=[
        CUDAExtension('gaussianpro',
            include_dirs=include_dirs,
            library_dirs=library_dirs,  
            libraries=['opencv_core', 'opencv_imgproc', 'opencv_highgui', 'opencv_imgcodecs'],  
            sources=[
                'PatchMatch.cpp', 
                'Propagation.cu',
                'pro.cpp'
            ],
            extra_compile_args={
                'cxx': ['-O3'],
                'nvcc': ['-O3',
                    '-gencode=arch=compute_89,code=sm_89',
                ]
            }),
    ],
    cmdclass={ 'build_ext' : BuildExtension }
)
